#!/usr/bin/env python3

import os
import sys
import tempfile
import subprocess
import macholib.MachO

from io import BytesIO
from pathlib import Path
from sys import argv

from macholib.mach_o import (
    LC_VERSION_MIN_MACOSX,
    LC_DATA_IN_CODE,
)

input_path = Path(argv[1].strip())

# Extract i386 slice from FAT binary if needed
result = subprocess.run(["lipo", input_path, "-info"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
if "i386" not in result.stdout:
    print("Binary is not 32-bit, nothing to do")
    sys.exit(0)


if "Non-fat" not in result.stdout:
    is_fat_binary = True
    print("Extracting 32-bit slice from fat binary")
    slice_temp_file, slice_path = tempfile.mkstemp()
    slice_path = Path(slice_path)
    os.close(slice_temp_file)
    result = subprocess.run(["lipo", input_path, "-thin", "i386", "-output", slice_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    
    # Initialize parser
    original_file = BytesIO(slice_path.read_bytes())
    machFile = macholib.MachO.MachO(slice_path)
else:
    is_fat_binary = False
    original_file = BytesIO(input_path.read_bytes())
    machFile = macholib.MachO.MachO(input_path)

# Strip min version command
for cmd in machFile.headers[0].commands:
    if (cmd[0].cmd == LC_VERSION_MIN_MACOSX):
        machFile.headers[0].changedHeaderSizeBy(-cmd[0].cmdsize)
        machFile.headers[0].commands.remove(cmd)
        machFile.headers[0].header.ncmds -= 1
        print("Removed LC_VERSION_MIN_MACOSX")

# Strip data-in-code command if zero
for cmd in machFile.headers[0].commands:
    if (cmd[0].cmd == LC_DATA_IN_CODE):
        if (cmd[1].datasize == 0):
            machFile.headers[0].changedHeaderSizeBy(-cmd[0].cmdsize)
            machFile.headers[0].commands.remove(cmd)
            machFile.headers[0].header.ncmds -= 1
            print("Removed LC_DATA_IN_CODE")
        else:
            print("LC_DATA_IN_CODE data size is non-zero!")
            sys.exit(-1)

# Align symbol table
symCmd = machFile.headers[0].getSymbolTableCommand()

oldSymOff = symCmd.symoff
newSymOff = (oldSymOff + 3) & ~(3)
print("Old symbol table {} new symbol table {}".format(oldSymOff, newSymOff))
symOffDelta = newSymOff - oldSymOff

# Align string table
oldStrOff = symCmd.stroff
newStrOff = ((oldStrOff + symOffDelta) + 3) & ~(3)
print("Old string table {} new string table {}".format(oldStrOff, newStrOff))
strOffDelta = newStrOff - (oldStrOff + symOffDelta)

# Write header
symCmd.symoff = newSymOff
symCmd.stroff = newStrOff

if is_fat_binary:
    new_file = slice_path.open("wb")
else:
    new_file = input_path.open("wb")

machFile.headers[0].write(new_file)  # Write header to file
original_file.seek(new_file.tell())  # Seek past header

# Copy rest of file and pad accordingly
if oldSymOff > oldStrOff:
    new_file.write(original_file.read(oldStrOff - original_file.tell()))
    new_file.write(b"\0" * strOffDelta)
    new_file.write(original_file.read(oldSymOff - original_file.tell()))
    new_file.write(b"\0" * symOffDelta)
    new_file.write(original_file.read())
else:
    new_file.write(original_file.read(oldSymOff - original_file.tell()))
    new_file.write(b"\0" * symOffDelta)
    new_file.write(original_file.read(oldStrOff - original_file.tell()))
    new_file.write(b"\0" * strOffDelta)
    new_file.write(original_file.read())

print("Symbol table aligned for " + input_path.name)

if is_fat_binary:
    print("Replacing 32-bit slice in fat binary")
    result = subprocess.run(["lipo", input_path, "-replace", "i386", slice_path, "-output", input_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    slice_path.unlink()
